# from .ZedCamera.main import zedcamera
import numpy as np
import os
import sys

from torch import double
from data.base_dataset import BaseDataset
sys.path.append(os.path.abspath("."))
import pyzed.sl as sl
# import pclpyd as pcl

# import datautils

def defaultinit(opt):
    init = sl.InitParameters()
    if opt.sl_depthmode.upper() in sl.DEPTH_MODE.__dict__.keys():
        init.depth_mode=getattr(sl.DEPTH_MODE,opt.sl_depthmode.upper())
    if opt.sl_coordinate_units.upper() in sl.UNIT.__dict__.keys():
        init.coordinate_units=getattr(sl.UNIT,opt.sl_coordinate_units.upper())
    if opt.sl_resolution.upper() in sl.RESOLUTION.__dict__.keys():
        init.camera_resolution=getattr(sl.RESOLUTION,opt.sl_resolution.upper())
    if opt.sl_coordinate_system.upper() in sl.COORDINATE_SYSTEM.__dict__.keys():
        init.coordinate_system=getattr(sl.COORDINATE_SYSTEM,opt.sl_coordinate_system.upper())
    init.depth_minimum_distance = opt.sl_depth_minimum_distance
    init.depth_maximum_distance = opt.sl_depth_maximum_distance
    return init
def defaultruntimeparameters(opt):
    runtime_parameters =sl.RuntimeParameters()
    if opt.sl_sensing_mode.upper() in sl.SENSING_MODE.__dict__.keys():
        runtime_parameters.sensing_mode = getattr(sl.SENSING_MODE,opt.sl_sensing_mode.upper())
    runtime_parameters.confidence_threshold = opt.sl_confidence_threshold
    runtime_parameters.texture_confidence_threshold = opt.sl_texture_confidence_threshold
    return runtime_parameters
def defaultres(opt):
    res = sl.Resolution()
    res.width = opt.sl_res_width
    res.height = opt.sl_res_height
    return res
def defaulttr_np(opt):
    mirror_ref = sl.Transform()
    mirror_ref.set_translation(sl.Translation(2.75,4.0,0))
    return mirror_ref.m

# def set_datactr(datactr):
#     datactr.vis = CVis(grab_pcl_fuc=datactr.get_data,
#                         switch_model_func=datactr.switch_model,
#                         switch_dataset_func=datactr.switch_dataset,
#                         switch_pcl_func=datactr.switch_pcl,
#                         )
#     datactr.config.Set_datactr_func=None
#     datactr.config.Set_exec_func = datactr.vis.run_pcl_dsiplay_process
#     datactr.set()
    
#     datactr.curmodel=0
#     for i,dataset in enumerate(datactr.config.datatsetstype):
#         datactr.set_vis_arg_func.append(importlib.import_module('dataflow.'+dataset+'.main').set_vis_args)
#     datactr.set_vis_arg_func[datactr.curdataset](datactr)
    
# def set_vis_args(datactr):
#     datactr.vis.set(xyzrelist=[1,-1,1],_3d=True)
# def get_dataset(config):  
#     return zedcamera()
    # if modeltypes is not None:
    #     datactr.model = [None]
        
    # datactr.vis = CVis(npoints=50000,xyzrelist=[1,-1,1],ballradius=1,
    #                    grab_pcl_fuc=datactr.data.grab_camera_pcl,
    #                    switch_model_func=datactr.switch_model,
    #                    )

class zeddataset(BaseDataset):
    def modify_commandline_options(parser, flagdict=None):
        """Add new dataset-specific options, and rewrite default values for existing options.

        Parameters:
            parser          -- original option parser
            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.

        Returns:
            the modified parser.
        """
        parser.add_argument('--sl_depthmode', type=str, default='PERFORMANCE', help='It decides the camera mode')
        parser.add_argument('--sl_coordinate_units', type=str, default='METER', help='')
        parser.add_argument('--sl_resolution', type=str, default='HD720', help='')
        parser.add_argument('--sl_coordinate_system', type=str, default='RIGHT_HANDED_Y_UP', help='')
        parser.add_argument('--sl_depth_minimum_distance', type=float, default=0.15, help='')
        parser.add_argument('--sl_depth_maximum_distance', type=float, default=30, help='')
        parser.add_argument('--sl_sensing_mode', type=str, default='STANDARD', help='')
        parser.add_argument('--sl_confidence_threshold', type=float, default=100, help='')
        parser.add_argument('--sl_texture_confidence_threshold', type=float, default=100, help='')
        parser.add_argument('--sl_res_width', type=int, default=720, help='')
        parser.add_argument('--sl_res_height', type=int, default=404, help='')

        return parser
    def change_options(self,opt):
        opt.yreverse=-1
        return opt
    def __init__(self,opt):
                #  init=defaultinit(),
                #  runtime_parameters=defaultruntimeparameters(),
                #  res=defaultres(),
                #  tr_np=defaulttr_np()):
        self.zed = sl.Camera()
        if self.zed is not None:
            self.zed.close()
        # init.coordinate_units = sl.UNIT.METER
        self.init = defaultinit(opt)
        status = self.zed.open(self.init)
        if status != sl.ERROR_CODE.SUCCESS:
            print(repr(status))
            self.zed.close()
            exit()
        self.runtime_parameters = defaultruntimeparameters(opt)
        self.res=defaultres(opt)

        self.point_cloud = sl.Mat(self.res.width, self.res.height, sl.MAT_TYPE.F32_C4, sl.MEM.CPU)
        self.tr_np = defaulttr_np(opt)
        self.model = [None]
        self.curmodel= 0
        self.npoints = None
    
    def __len__(self):
        return 0
    def __getitem__(self,idx):
        if self.zed.grab(self.runtime_parameters) == sl.ERROR_CODE.SUCCESS:
            # vis.changed = True
            self.zed.retrieve_measure(self.point_cloud, sl.MEASURE.XYZRGBA, sl.MEM.CPU)
            calibration_params = self.zed.get_camera_information().calibration_parameters
            fx = calibration_params.left_cam.fx
            fy = calibration_params.left_cam.fy
            cx = calibration_params.left_cam.cx
            cy = calibration_params.left_cam.cy
            # mtx = np.array([[fx,0,cx],[0,fy,cy],[0,0,1]])
            # disto = calibration_params.left_cam.disto

            cleanpointcloud = self.point_cloud.get_data().astype(np.float32)
            cleanpointcloud.dot(self.tr_np)
            cleanpointcloud = cleanpointcloud.reshape(cleanpointcloud.shape[0]*cleanpointcloud.shape[1],4)
            cleanpointcloud = cleanpointcloud[~np.isnan(cleanpointcloud).any(axis=1)]
            cleanpointcloud = cleanpointcloud[~np.isinf(cleanpointcloud).any(axis=1)]
            # cleanpointcloud[:,3] = ((-1 << 31) - (cleanpointcloud[:,3].astype(np.int64) >> 32 )).astype(np.float32)
            # cleanpointcloud = cleanpointcloud[~np.isnan(cleanpointcloud).any(axis=1)]
            args=np.array([6]).astype(np.float64)
            pointcloud = pcl.endecode_color(cleanpointcloud,args).astype(np.float32)
            # pointcloud = cv2.undistort()

            f = (fx+fy) /2
            x = pointcloud[:,0] - cx
            y = pointcloud[:,1] - cy
            cos_theta = f/np.sqrt(x**2 + y**2 + f ** 2)
            pointcloud[:,2] *= cos_theta
            pointcloud[:,3:6] = np.array([pointcloud[:,5]*10,pointcloud[:,4],pointcloud[:,3]*10]).T
            return pointcloud,True